Conversation
Summary of ChangesHello @meichangsu1, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request resolves an inaccuracy in loss metric calculation when Sequence Parallelism is active, particularly when combined with FSDP. The core solution involves introducing a compensation mechanism for FSDP's averaging behavior and refining the loss reduction process through custom autograd functions. These changes ensure that gradients are scaled correctly and the reported loss accurately reflects the model's performance under distributed training configurations. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request fixes an issue with loss and gradient scaling when sequence parallelism (SP) is enabled with FSDP/DDP. The changes introduce a compensation factor in the backward pass to correct for the extra averaging performed by FSDP across SP ranks. This is controlled by a new compensate_fsdp_avg flag. The implementation of loss reduction has been refactored to use torch.autograd.Function for cleaner custom gradient computation.
My review identifies a critical issue in the new implementation of _ReduceSequenceParallelSum which could lead to incorrect gradient scaling, and a medium-severity issue in the README.md documentation. I've provided detailed comments and suggested fixes for both.
…d functions Replace manual gradient handling with `torch.autograd.Function` subclasses `_ReduceSequenceParallelLoss` and `_ReduceSequenceParallelSum` to compute global loss via autograd-aware all-reduce. This simplifies the logic for both sum and mean reductions, improves gradient correctness, and removes the need for separate metric scaling when `world_size > 1`.
Add `compensate_fsdp_avg` config flag to adjust loss reduction when sequence parallel (SP) is combined with FSDP or accelerate DDP/FSDP. This prevents gradient magnitude from being incorrectly scaled down by an extra factor of SP world size during data-parallel averaging. - In `GatherLoss` backward, scale gradients by SP world size before splitting, so downstream FSDP averaging does not shrink this path. - In `SequenceParallelStrategy.reduce_loss`, apply a compensation factor (ulysses_size) when `compensate_fsdp_avg` is enabled. - Automatically set `compensate_fsdp_avg=True` in `TransformersModel` when using NativeFSDPStrategy or AccelerateStrategy with both SP and data parallelism active.
meichangsu1
force-pushed
the
seq_unitest_ljl_fix
branch
from
f5ca919 to
9d239da
Compare
- Add 'kernels' as an optional dependency group in pyproject.toml - Refactor CI container test script to use a reusable installation function - Install twinkle with kernels in both debug and release modes for consistency - Improve maintainability by centralizing the installation command
Update `_load_from_hub` function to handle API changes in `select_revision_or_version` and `get_kernel` calls. The changes introduce try-except blocks to catch `TypeError` exceptions, allowing the function to work with both modern keyword-based APIs and older positional argument variants. This ensures compatibility across different versions of the kernels module without breaking existing functionality.
1 participant
No description provided.